import json
import os
import cv2
import numpy as np
from tqdm import tqdm

def convert1():
    jsonPath = "Z:\\电科未检测到驾驶员视频附件_拆帧\\json"
    txtPath = "Z:\\电科未检测到驾驶员视频附件_拆帧\\label"
    imgPath = "Z:\\电科未检测到驾驶员视频附件_拆帧\\image_label"
    for jsonname in tqdm(os.listdir(jsonPath)):
        jsonfile = os.path.join(jsonPath, jsonname)
        txtfile = os.path.join(txtPath, jsonname.split(".")[0] + ".txt")
        imagefile = os.path.join(imgPath, jsonname.split(".")[0] + ".jpg")
        if not os.path.exists(imagefile):
            continue
        if os.path.exists(txtfile):
            continue
        image = cv2.imdecode(np.fromfile(imagefile, dtype=np.uint8), -1)
        height, width = image.shape[:2]
        if not os.path.exists(txtPath):
            os.makedirs(txtPath)
        with open(jsonfile, 'r', encoding='utf8')as fp:
            json_data = json.load(fp)
            # print(json_data)
            num_obj = len(json_data["shapes"])
            if num_obj == 0:
                continue
            with open(txtfile, "w", encoding='utf8') as ft:
                for object in json_data["shapes"]:
                    if (object["label"] == "face"):
                        label = 0
                    elif (object["label"] == "hand"):
                        label = 1
                    elif (object["label"] == "cigarette"):
                        label = 2
                    elif (object["label"] == "cellphone"):
                        label = 3
                    # 增加标注不为为左上和右下形式的保护
                    x1 = min(object["points"][0][0], object["points"][1][0])
                    y1 = min(object["points"][0][1], object["points"][1][1])
                    x2 = max(object["points"][0][0], object["points"][1][0])
                    y2 = max(object["points"][0][1], object["points"][1][1])
                    object_x_center = (int(x1) + int(x2)) / 2.0 / float(width)
                    object_y_center = (int(y1) + int(y2)) / 2.0 / float(height)
                    object_width = (int(x2) - int(x1)) / float(width)
                    object_height = (int(y2) - int(y1)) / float(height)
                    ft.write(str(label) + " " + str(object_x_center) + " " + str(object_y_center) + " " + str(
                        object_width) + " " + str(object_height) + "\n")

# 只转换包换安全带的标注
def convert2():
    jsonPath = "D:\\dataset\\belt_detect\\train\\json24-40"
    txtPath = "D:\\dataset\\belt_detect\\train\\label1-40"
    imgPath = "D:\\dataset\\DSM_Dataset_class4_20220211_fukang\\train\\image"
    for jsonname in tqdm(os.listdir(jsonPath)):
        jsonfile = os.path.join(jsonPath, jsonname)
        txtfile = os.path.join(txtPath, jsonname.split(".")[0] + ".txt")
        imagefile = os.path.join(imgPath, jsonname.split(".")[0] + ".jpg")
        image = cv2.imdecode(np.fromfile(imagefile, dtype=np.uint8), -1)
        height, width = image.shape[:2]
        if not os.path.exists(txtPath):
            os.makedirs(txtPath)
        with open(jsonfile, 'r', encoding='utf8')as fp:
            json_data = json.load(fp)
            # print(json_data)
            with open(txtfile, "w", encoding='utf8') as ft:
                for object in json_data["shapes"]:
                    if (object["label"] == "belt"):
                        label = 0
                    else:
                        continue
                    # 增加标注不为为左上和右下形式的保护
                    x1 = min(object["points"][0][0], object["points"][1][0])
                    y1 = min(object["points"][0][1], object["points"][1][1])
                    x2 = max(object["points"][0][0], object["points"][1][0])
                    y2 = max(object["points"][0][1], object["points"][1][1])
                    object_x_center = (int(x1) + int(x2)) / 2.0 / float(width)
                    object_y_center = (int(y1) + int(y2)) / 2.0 / float(height)
                    object_width = (int(x2) - int(x1)) / float(width)
                    object_height = (int(y2) - int(y1)) / float(height)
                    ft.write(str(label) + " " + str(object_x_center) + " " + str(object_y_center) + " " + str(
                        object_width) + " " + str(object_height) + "\n")

# 从val.json中获取json名称
def convert3():
    jsonFile = "D:\\dataset\\DSM_Dataset_class4_20220211_fukang_v1_resample3\\train\\train_v1_resample3_xinQ37133_2136.json"
    jsonPath = "D:\\dataset\\DSM_Dataset_class4_20220211_fukang_v1_resample3\\train\\json"
    txtPath = "D:\\dataset\\DSM_Dataset_class4_20220211_fukang_v1_resample3\\train\\label_v1_resample3_xinQ37133"
    imgPath = "D:\\dataset\\DSM_Dataset_class4_20220211_fukang_v1_resample3\\train\\image"
    with open(jsonFile, "r", encoding="utf-8") as f:
        jsonContent = json.load(f)
        for img_list in tqdm(jsonContent["images"]):
            jsonname = img_list["file_name"].split(".")[0] + ".json"
            jsonfile = os.path.join(jsonPath, jsonname)
            txtfile = os.path.join(txtPath, jsonname.split(".")[0] + ".txt")
            imagefile = os.path.join(imgPath, jsonname.split(".")[0] + ".jpg")
            if not os.path.exists(imagefile):
                continue
            if os.path.exists(txtfile):
                continue
            image = cv2.imdecode(np.fromfile(imagefile, dtype=np.uint8), -1)
            height, width = image.shape[:2]
            if not os.path.exists(txtPath):
                os.makedirs(txtPath)
            with open(jsonfile, 'r', encoding='utf8')as fp:
                json_data = json.load(fp)
                # print(json_data)
                with open(txtfile, "w", encoding='utf8') as ft:
                    for object in json_data["shapes"]:
                        if (object["label"] == "face"):
                            label = 0
                        elif (object["label"] == "hand"):
                            label = 1
                        elif (object["label"] == "cigarette"):
                            label = 2
                        elif (object["label"] == "cellphone"):
                            label = 3
                        # 增加标注不为为左上和右下形式的保护
                        x1 = min(object["points"][0][0], object["points"][1][0])
                        y1 = min(object["points"][0][1], object["points"][1][1])
                        x2 = max(object["points"][0][0], object["points"][1][0])
                        y2 = max(object["points"][0][1], object["points"][1][1])
                        object_x_center = (int(x1) + int(x2)) / 2.0 / float(width)
                        object_y_center = (int(y1) + int(y2)) / 2.0 / float(height)
                        object_width = (int(x2) - int(x1)) / float(width)
                        object_height = (int(y2) - int(y1)) / float(height)
                        ft.write(str(label) + " " + str(object_x_center) + " " + str(object_y_center) + " " + str(
                            object_width) + " " + str(object_height) + "\n")

def delete_align():
    txt_path = "D:\\dataset\\belt_detect\\val\\label"
    img_path = "D:\\dataset\\belt_detect\\val\\image"
    for imgname in os.listdir(img_path):
        imgfile = os.path.join(img_path, imgname)
        txtName = imgname.split(".")[0] + ".txt"
        txtfile = os.path.join(txt_path, txtName)
        if not os.path.exists(txtfile):
            print("remove " + imgfile)
            os.remove(imgfile)



# 将labelme标注文件转化为yolov5所需数据格式
if __name__ == "__main__":
    convert1()

